package org.znxs.znagent_s.rag;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;

/**
 * RAG 检索增强 自定义属性 Advisor 工厂
 * <p>
 * 可选功能，自定义文档过滤检索器、自定义查询增强器（例如空上下文处理器）、自定义兜底策略Advisor等
 * <p>
 * 返回都是Advisor
 */
@Slf4j
public class LoveAppRagCustomAdvisorFactory {

    /**
     * 创建自定义 Advisor
     *
     * @param vectorStore 向量存储
     * @param status      状态
     * @return Advisor 顾问
     */
    public static Advisor createRagCustomAdvisor(VectorStore vectorStore, String status) {
        Filter.Expression expression = new FilterExpressionBuilder()
                .eq("status", status)
                .build();

        DocumentRetriever documentRetriever = VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .filterExpression(expression) // 只有满足这个过滤表达式的才能检索通过
                .similarityThreshold(0.5) // 相似度过滤 只有达到这个相似度才能检索通过
                .topK(3) // 检索前几个
                .build();

//        // 通过官方自带的检索增强顾问来 构建实例
//        return RetrievalAugmentationAdvisor.builder()
//                // 使用自定义文档检索处理
//                .documentRetriever(documentRetriever)
//                // 使用查询增强处理 空上下文处理
//                .queryAugmenter(LoveAppContextualQueryAugmenterFactory.createInstance())
//                .build();
        // 这里RetryableDocumentRetriever的参数为：文档过滤检索实例，未符合重新检索次数，向量存储对象
        RetryableDocumentRetriever retryableDocumentRetriever = new RetryableDocumentRetriever(documentRetriever, 2, vectorStore);
        return RetrievalAugmentationAdvisor.builder()
                // 使用自定义文档检索处理
                .documentRetriever(retryableDocumentRetriever)
                // 使用增强处理 空文档处理
                .queryAugmenter(LoveAppContextualQueryAugmenterFactory.createInstance())
                .build();

    }
}
